#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/uio.h>
#include <sys/socket.h>

/* in the <sys/socket.h>, but not include in??? */
struct ucred {
    uid_t uid;
    gid_t gid;
    pid_t pid;
};

/* size of control buffer to send/recv one file descriptor */
#define MAXLINE         8192
#define CREDSTRUCT      ucred
#define CR_UID          uid
#define CREDOPT         SO_PASSCRED
#define SCM_CREDTYPE    SCM_CREDENTIALS

#define RIGHTSLEN       CMSG_LEN(sizeof(int))
#define CREDSLEN        CMSG_LEN(sizeof(struct CREDSTRUCT))
#define CONTROLLEN      (RIGHTSLEN + CREDSLEN)

static struct cmsghdr *cmptr = NULL;    /* malloc'ed first time */

/*
 * Pass a file descriptor to another process.
 * If fd < 0, then -fd is send back instead as the error status.
 */
int send_fd(int fd, int fd_to_send)
{
    struct CREDSTRUCT   *credp;
    struct cmsghdr      *cmp;
    struct iovec        iov[1];
    struct msghdr       msg;
    char                buf[2]; /* send_fd()/recv_fd() 2-byte protocol */

    iov[0].iov_base = buf;
    iov[0].iov_len = 2;
    msg.msg_iov = iov;
    msg.msg_iovlen = 1;
    msg.msg_name = NULL;
    msg.msg_namelen = 0;
    if (fd_to_send < 0) {
        msg.msg_control = NULL;
        msg.msg_controllen = 0;
        buf[1] = -fd_to_send;   /* nonzero status means error*/
        if (buf[1] == 0)        /* -256 */
            buf[1] = 1;
    }
    else {
        if (cmptr == NULL && (cmptr = malloc(CONTROLLEN)) == NULL)
            return -1;
        msg.msg_control = cmptr;
        msg.msg_controllen = CONTROLLEN;
        cmp = cmptr;
        cmp->cmsg_level = SOL_SOCKET;
        cmp->cmsg_type = SCM_RIGHTS;
        cmp->cmsg_len = RIGHTSLEN;
        *(int *)CMSG_DATA(cmp) = fd_to_send;
        
        cmp = CMSG_NXTHDR(&msg, cmp);
        cmp->cmsg_level = SOL_SOCKET;
        cmp->cmsg_type = SCM_CREDTYPE;
        cmp->cmsg_len = CREDSLEN;
        credp = (struct CREDSTRUCT *)CMSG_DATA(cmp);
        credp->uid = geteuid();
        credp->gid = getegid();
        credp->pid = getpid();

        buf[1] = 0;             /* zero status means OK */
    }
    buf[0] = 0;
    if (sendmsg(fd, &msg, 0) != 2)
        return -1;

    return 0;
}

/*
 * Receive a file descriptor from a server process. Also, any data received 
 * is pass to (*userfunc)(STDERR_FILENO, buf, nbytes).
 * We have a 2-byte protocol for receiving the fd from send_fd().
 */
int recv_fd(int fd, uid_t *uidptr, 
    ssize_t (*userfunc)(int, const void *, size_t)) 
{
    struct cmsghdr      *cmp;
    struct CREDSTRUCT   *credp;
    int                 newfd, nr, status;
    char                *ptr;
    char                buf[MAXLINE];
    struct iovec        iov[1];
    struct msghdr       msg;
    const int           on = 1;

    status = -1;
    newfd = -1;
    if (setsockopt(fd, SOL_SOCKET, CREDOPT, &on, sizeof(int)) < 0) {
        perror("setsockopt error");
        return -1;
    }

    for (;;) {
        iov[0].iov_base = buf;
        iov[0].iov_len = sizeof(buf);
        msg.msg_iov = iov;
        msg.msg_iovlen = 1;
        msg.msg_name = NULL;
        msg.msg_namelen = 0;
        if (cmptr == NULL && (cmptr = malloc(CONTROLLEN)) == NULL)
            return -1;
        msg.msg_control = cmptr;
        msg.msg_controllen = CONTROLLEN;
        if ((nr = recvmsg(fd, &msg, 0)) < 0) {
            perror("recvmsg error");
            exit(0);
        }
        else if (nr == 0) {
            perror("connection closed bt server");
            return -1;
        }
        
        /*
         * See if this is the final data with null & status. NULL is next
         * to last byte of buffer,status byte is last byte.
         * Zero status means there is a file descriptor to receive.
         */
        for (ptr = buf; ptr < &buf[nr];) {
            if (*ptr++ == 0) {
                if (ptr != &buf[nr-1]) {
                    printf("message format error\n");
                    exit(0);
                }
                status = *ptr & 0xFF;   /* prevent signal extension */
                if (status == 0) {
                    if (msg.msg_controllen != CONTROLLEN) {
                        printf("status = 0 but no fd\n");
                        exit(0);
                    }
                    /* process control data */
                    for (cmp = CMSG_FIRSTHDR(&msg); cmp != NULL; 
                        cmp = CMSG_NXTHDR(&msg, cmp)) {
                        if (cmp->cmsg_level != SOL_SOCKET)
                            continue;
                        switch (cmp->cmsg_type) {
                        case SCM_RIGHTS:
                            newfd = *(int *)CMSG_DATA(cmp);
                            break;
                        case SCM_CREDTYPE:
                            credp = (struct CREDSTRUCT *)CMSG_DATA(cmp);
                            *uidptr = credp->CR_UID;
                        }
                    }
                }
                else {
                    newfd = -status;
                }
                nr -= 2;
            }
        }
        if (nr > 0 && (*userfunc)(STDERR_FILENO, buf, nr) != nr)
            return -1;
        if (status >= 0)
            return newfd;
    }
}

/*
 * Used when we had planned to send an fd using send_fd(), but encountered 
 * an error instead. We send the error back using the send_fd()/recv_fd()
 * protocol.
 */
int send_err(int fd, int errcode, const char *msg)
{
    int n;

    if ((n = strlen(msg)) > 0) 
        if (write(fd, msg, n) != n)
            return -1;

    if (errcode >= 0)
        errcode = -1;   /* must be negative */

    if (send_fd(fd, errcode) < 0)
        return -1;

    return 0;
}
